from functools import partial

import numpy as np
import torch

import sko
from scene.atm_opt.model.ga_atm import GA_ATM, parallel_atm_costfunc, gpu_atm_costfunc


class ATM_OPT:
    PROJECT_ROOT = ""
    DATA_FOLDER_PATH = ""
    SCRIPTS_ROOT = ""
    DATA = ""

    def __init__(self, args, DATA):
        self.args = args

        self.loadData(DATA)

    # 加载数据
    def loadData(self, DATA):
        '''
            参数说明：
            numATM n : 珠海东区ATM机的个数
            computeIteration k : 总成本计算周期
            D_1: 由珠海工行xgboost函数预测得到的未来k天各ATM机取款预测值
            D_2: 由珠海工行xgboost函数预测得到的未来k天各ATM机存款预测值
            distance_list: 辖区各ATM机距离表
            c_label: 各ATM机按现金使用率的高低划分的类型表
            aver_1:根据历史数据统计得到的各ATM机日均净取款表
        '''
        self.DATA = DATA
        self.DATA.C = 0
        self.D = self.DATA.D
        # self.distance_list = self.DATA.distance_list

        self.CUTOFF_LIST = self.DATA.cutoff_list.flatten()
        assert np.all(self.CUTOFF_LIST >= 0)
        self.aver_1 = self.DATA.aver_1.flatten()

        '''
        寻找通过均值和方差对D2存款值进行估计的最优参数
        '''
        self.DATA.C = 0

    # 应用ALNS算法寻找最优的加钞序号、加钞值，最优成本
    def applyALNS(self):
        costfun = partial(
            parallel_atm_costfunc,
            DATA=self.DATA,
            args=self.args,
            device=self.args.device,
        )
        sko.tools.set_run_mode(costfun, 'vectorization')
        self.ga = GA_ATM(costfun, self.DATA, self.args)
        self.ga.to(self.args.device)
        a, b = self.ga.run()
        DATA = self.DATA
        args = self.args

        (_, self.I, self.cutoff_list,
         (self.over_cutoff, self.usage_rate, self.short_rate, self.add_freq)) = gpu_atm_costfunc(
            y=torch.stack([a]),
            D=torch.from_numpy(DATA.D.astype(np.int32)).float().to(args.device),
            aver_1=torch.from_numpy(DATA.aver_1).int().to(args.device),
            max_cap=torch.from_numpy(DATA.max_cap.astype(np.int32)).int().to(args.device),
            w_0=torch.from_numpy(DATA.w_0).int().to(args.device),
            cutoff_list=torch.from_numpy(DATA.cutoff_list).int().to(args.device),
            I_0=torch.from_numpy(DATA.I_0).int().to(args.device),
            avg_D2=torch.from_numpy(DATA.avg_D2).int().to(args.device),
            C=DATA.C,
            std_D2=torch.from_numpy(DATA.std_D2).float().to(args.device),
            cutoff_limit=torch.from_numpy(DATA.cutoff_limit).int().to(args.device),
            args=args
        )
        self.w = a.reshape(1, -1, DATA.w_0.size)

    # 最优加钞金额
    def getBestAddMount(self):
        nxt = 0
        while (self.args.day - 1 + nxt) % 7 + 1 in self.args.restday:
            nxt += 1
        return self.w[:, nxt] / (1 << self.args.bit_precision)

    # 最优的总成本对应的缺钞率均值
    def getBestPenaltyMean(self):
        return self.short_rate

    # 最优的总成本对应的现金使用率均值
    def getBestCashUsageMean(self):
        return self.usage_rate

    # 最优的总成本对应的加钞频率均值
    def getBestCashAddFreq(self):
        return self.add_freq

    # 给出第1-7日的出车线路选择，0号代表现金中心
    def getBestRoute(self):
        return 0

    # 更新第1天加钞后的清机系数
    def getATMClearRate(self):
        nxt = 0
        while (self.args.day - 1 + nxt + 1) % 7 + 1 in self.args.restday:
            nxt += 1
        return self.cutoff_list[:, nxt + 1].int()
